# encoding=utf-8
import torch
from networks import get_norm_layer, ResUNet

# load model
device = torch.device("cpu")
model = ResUNet(3, 3, 64, norm_layer=get_norm_layer())
model_parameter = torch.load("120_net_G_A.pth", map_location=device)
model.load_state_dict(model_parameter)
model.eval()


# convert by onnx
torch.onnx.export(
    model,
    torch.randn(1, 3, 512, 512),
    "export_funds_stillgan.onnx",
    verbose=True,
    opset_version=16,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output']
)

# test onnx
import onnxruntime
session = onnxruntime.InferenceSession("export_funds_stillgan.onnx")
inputs = {"input":  torch.randn(1, 3, 512, 512).numpy()}
out = session.run(None, inputs)
print(out)








